use image::{ImageBuffer, RgbImage};
use linfa::{
    linalg::{Norm, SVD},
    prelude::*,
};
use mnist::{Mnist, MnistBuilder};
use ndarray::{prelude::*, ArcArray2};
use ndarray_rand::{rand_distr::Uniform, RandomExt};
use plotters::prelude::*;
use show_image::{event, ImageInfo, ImageView, WindowOptions};

fn main() {
    println!("nihao");
    pca1();
    //mnist();
    //handle_mnist();
}

// 使用梯度上升发求解pca
fn pca1() {
    let mut x: Array2<f64> = Array2::zeros((100, 0));
    let x1 = Array1::random(100, Uniform::new(0.0, 100.0));
    x.push_column(x1.view()).unwrap();
    let x2 = x1 * 0.75 + 3.0 + Array1::random(100, Uniform::new(0.0, 20.0));
    x.push_column(x2.view()).unwrap();
    println!("{}", x);
    // 显示图形
    let root = SVGBackend::new("./examples/pca/pca.svg", (600, 400)).into_drawing_area();
    root.fill(&WHITE).unwrap();
    let mut chat = ChartBuilder::on(&root)
        .set_left_and_bottom_label_area_size(40)
        .build_cartesian_2d(-110.0..110.0, -110.0..110.0)
        .unwrap();

    chat.configure_mesh()
        .disable_x_mesh()
        .disable_y_mesh()
        .draw()
        .unwrap();

    let x1 = x.column(0);
    let x2 = x.column(1);
    let r = x1.iter().zip(x2);
    chat.draw_series(r.map(|(x, y)| Circle::new((*x, *y), 4, BLUE.filled())))
        .unwrap();
    let x = mean_to_zero(x.view());
    chat.draw_series(
        x.column(0)
            .iter()
            .zip(x.column(1))
            .map(|(x, y)| Circle::new((*x, *y), 4, RED.filled())),
    )
    .unwrap();

    println!("{}", x.mean_axis(Axis(0)).unwrap());
    println!("{}", x.mean_axis(Axis(1)).unwrap());
    println!("{}", x.column(0));
    let wawa: Array2<f64> = array![[1.0], [3.0]];
    let haha: Array2<f64> = array![[5.0], [2.0]];
    println!("{}", wawa.r2(&haha).unwrap());
    println!(
        "{}",
        wawa.mean_squared_error(&haha)
            .unwrap()
            .get(0)
            .unwrap()
            .sqrt()
    );
    println!("{}", wawa.mean_absolute_error(&haha).unwrap());
    println!("{}", wawa.mean_squared_log_error(&haha).unwrap());
    println!("{}", wawa.mean_absolute_percentage_error(&haha).unwrap());
    let initial_w = Array2::random((2, 1), Uniform::new(0.0, 1.0));
    println!("{}", initial_w);
    let eta = 0.001;
    let res = first_component(df_debug, x.view(), initial_w.view(), eta, 10000, 1e-8);
    println!("res: {}", res);
    // 第一主成分
    let res = first_component(df_math, x.view(), initial_w.view(), eta, 10000, 1e-8);
    println!("res: {}", res);
    let points: Vec<(f64, f64)> = vec![(0.0, 0.0), (res.row(0)[0] * 30.0, res.row(1)[0] * 30.0)];
    chat.draw_series(LineSeries::new(points, &YELLOW)).unwrap();

    // 求出第二主成分
    let w = res.to_shared().reshape(res.len_of(Axis(0)));
    let mut x2: Array2<f64> = Array::zeros((0, x.len_of(Axis(1))));
    for i in 0..x.len_of(Axis(0)) {
        let row = x.row(i);
        x2.push_row((&row - row.dot(&w) * &w).view()).unwrap();
    }
    chat.draw_series(
        x2.column(0)
            .iter()
            .zip(x2.column(1))
            .map(|(x, y)| Circle::new((*x, *y), 1, BLACK.filled())),
    )
    .unwrap();
    let res2 = first_component(df_math, x2.view(), initial_w.view(), eta, 10000, 1e-8);
    println!("res2: {}", res2);
    chat.draw_series(LineSeries::new(
        vec![
            (0.0, 0.0),
            (
                *res2.get((0, 0)).unwrap() * 30.0,
                *res2.get((1, 0)).unwrap() * 30.0,
            ),
        ],
        &BLUE,
    ))
    .unwrap();
    println!("{}", res.t().dot(&res2));

    let all = first_n_component(2, x.view(), 0.001, 10000);
    println!("all: {}", all);

    // 使用我们自己封装的库进行调用
    let mut pca = Pca::new(1, x.view());
    pca.fit(0.001, 10000);
    println!("components: {}", pca.get_components());
    // println!("x_reduction: {}", pca.transform());

    // 恢复原来的数据并显示出来
    let restore = pca.invers_tranform(pca.transform().view());
    //println!("x_restore: {}", restore);
    chat.draw_series(
        restore
            .column(0)
            .iter()
            .zip(restore.column(1))
            .map(|(&x, &y)| Circle::new((x, y), 4, MAGENTA.filled())),
    )
    .unwrap();

    // 使用linfa中的pca进行调用
    let x_row_count = x.len_of(Axis(0));
    let dataset: Dataset<f64, f64, Ix1> = Dataset::new(x.clone(), Array1::zeros(x_row_count));
    let pca = linfa_reduction::Pca::params(2).fit(&dataset).unwrap();
    let _x_reduction = pca.predict(dataset.records());
    //println!("x_reduction: {}", x_reduction);
    println!("pca: {:?}", pca);
    println!("奇异值: {}", pca.singular_values());
    println!("{}", pca.explained_variance());
    println!("{}", pca.explained_variance_ratio());

    let (_s, v, _d) = x.svd(false, false).unwrap();
    println!("v: {}", v);
}

fn handle_mnist() {
    let (trn_size, _rows, _cols) = (50_000, 28, 28);
    let Mnist {
        trn_img,
        trn_lbl,
        tst_img,
        tst_lbl,
        ..
    } = MnistBuilder::new()
        .base_path("data/")
        .label_format_digit()
        .training_set_length(trn_size)
        .validation_set_length(10_000)
        .test_set_length(10_000)
        .download_and_extract()
        .finalize();
    println!("{}", trn_img.len());

    let x_train = Array2::from_shape_vec((50_000, 784), trn_img)
        .expect("Error converting images to Array3 struct")
        .mapv(|x| x as f64);
    let y_train = Array1::from_vec(trn_lbl).mapv(|x| x as f64);

    let x_test = Array2::from_shape_vec((10_000, 784), tst_img)
        .unwrap()
        .mapv(|x| x as f64);
    let y_test = Array1::from_vec(tst_lbl).mapv(|x| x as f64);

    println!("x_train shape: {:?}", x_train.shape());
    println!("x_train shape: {:?}", y_train.shape());
    println!("x_test shape: {:?}", x_test.shape());
    println!("y_test shape: {:?}", y_test.shape());
    let train_dataset = Dataset::new(x_train, y_train);
    let _test_dataset = Dataset::new(x_test, y_test);

    // 降维
    let pca_p = linfa_reduction::Pca::params(5);
    let pca = pca_p.fit(&train_dataset).unwrap();
    let train_dataset = pca.transform(train_dataset);
    println!(
        "train_dataset recoreds shape: {:?}",
        train_dataset.records().shape()
    );

    // let k = Knn::fit(&x_train, &y_train, linfa_nn::KdTree);
    // let ret = k.predict(&x_test, 5, linfa_nn::distance::L2Dist).unwrap();
    // let score = ret.mean_absolute_percentage_error(&y_test).unwrap();
    // println!("score: {}", score);
}
// mnist手写识别的数据
fn mnist() {
    let (trn_size, _rows, _cols) = (50_000, 28, 28);
    let Mnist {
        trn_img, trn_lbl, ..
    } = MnistBuilder::new()
        .base_path("data/")
        .label_format_digit()
        .training_set_length(trn_size)
        .validation_set_length(10_000)
        .test_set_length(10_000)
        .download_and_extract()
        .finalize();

    let item_num = 7;
    return_item_description_from_number(trn_lbl[item_num]);

    let train_data = Array3::from_shape_vec((50_000, 28, 28), trn_img)
        .expect("Error converting images to Array3 struct")
        .mapv(|x| x as f64 / 256.);

    let image = bw_ndarray2_to_rgb_image(train_data.slice(s![item_num, .., ..]).to_owned());
    let window_options = WindowOptions::new();
    let window_options = window_options
        .set_size([100, 100])
        .set_resizable(true)
        .set_preserve_aspect_ratio(true)
        .set_fullscreen(false);

    show_image::run_context(move || {
        let show_iamge_data = image.to_vec();
        let image = ImageView::new(ImageInfo::rgb8(28, 28), &show_iamge_data);
        let window = show_image::create_window("image", window_options).unwrap();
        window.set_image("wawa", image).unwrap();

        // event
        for e in window.event_channel().unwrap() {
            if let event::WindowEvent::KeyboardInput(event) = e {
                if event.input.key_code == Some(event::VirtualKeyCode::Escape) {
                    println!("退出循环");
                    break;
                }
            }
        }

        Result::<(), Box<dyn std::error::Error>>::Ok(())
    });
}

// 返回对应label的数字
fn return_item_description_from_number(val: u8) {
    let description = match val {
        0 => "0",
        1 => "1",
        2 => "2",
        3 => "3",
        4 => "4",
        5 => "5",
        6 => "6",
        7 => "7",
        8 => "8",
        9 => "9",
        _ => panic!("An unrecognized label was used..."),
    };
    println!(
        "Based on the '{}' label, this image should be a : {}",
        val, description
    );
    println!("Hit [ESC] to exit....");
}

// 生成image数据
fn bw_ndarray2_to_rgb_image(arr: Array2<f64>) -> RgbImage {
    assert!(arr.is_standard_layout());
    let (width, height) = (arr.ncols(), arr.ncols());
    let mut img: RgbImage = ImageBuffer::new(width as u32, height as u32);
    for y in 0..height {
        for x in 0..width {
            let val = (arr[[y, x]] * 255.) as u8;
            img.put_pixel(x as u32, y as u32, image::Rgb([val, val, val]))
        }
    }
    img
}
struct Pca {
    n: usize,
    x: ArcArray2<f64>,
    pca: Array2<f64>,
}

impl Pca {
    fn new(n_componenets: usize, x: ArrayView2<f64>) -> Self {
        Self {
            n: n_componenets,
            x: x.to_shared(),
            pca: Array2::zeros((0, 0)),
        }
    }

    fn fit(&mut self, eta: f64, n_iters: usize) {
        let column_count = self.x.len_of(Axis(1));
        assert!(self.n >= 1 && self.n <= column_count);
        let res = Self::first_n_component(self.n, self.x.view(), eta, n_iters);
        self.pca = res;
    }

    // 降维度
    fn transform(&self) -> Array2<f64> {
        self.x.dot(&self.pca)
    }

    // 把降维的数据恢复成原来的数据， 需要知道主成分+降维后的数据，就能恢复成原来数据
    fn invers_tranform(&self, x: ArrayView2<f64>) -> Array2<f64> {
        let w = self.pca.t();
        assert!(x.len_of(Axis(1)) == w.len_of(Axis(0)));
        x.dot(&w)
    }

    fn get_components(&self) -> Array2<f64> {
        self.pca.clone()
    }

    // 均值归一化
    // 返回一个新的矩阵
    fn mean_to_zero(x: ArrayView2<f64>) -> Array2<f64> {
        &x - x.mean_axis(Axis(0)).unwrap()
    }

    // 梯度上升法求解PCA
    // 损失函数:f(x) = 1/m sigma(i=1,m) (XW)^2
    // 梯度是: 2/m * X^T * (XW)

    // 我们的目标函数
    fn f(x: ArrayView2<f64>, w: ArrayView2<f64>) -> f64 {
        // 生成一个列向量
        let c1 = x.dot(&w);
        let c1 = &c1 * &c1;
        c1.sum() / x.len_of(Axis(0)) as f64
    }

    // 目标函数对应梯度的方法
    fn df_math(x: ArrayView2<f64>, w: ArrayView2<f64>) -> Array2<f64> {
        x.t().dot(&x.dot(&w)) * 2f64 / x.len_of(Axis(0)) as f64
    }

    // debug函数
    fn df_debug(x: ArrayView2<f64>, w: ArrayView2<f64>) -> Array2<f64> {
        let mut res: Array2<f64> = Array2::zeros((w.len_of(Axis(0)), 1));
        for i in 0..w.len_of(Axis(0)) {
            let mut w_1 = w.to_owned();
            w_1.row_mut(i)[0] += 1e-8;
            let mut w_2 = w.to_owned();
            w_2.row_mut(i)[0] -= 1e-8;
            res.row_mut(i)[0] = (Self::f(x, w_1.view()) - Self::f(x, w_2.view())) / (2.0 * 1e-8);
        }
        res
    }

    // 梯度上升法求出第一主成分
    fn first_component(
        df: impl Fn(ArrayView2<f64>, ArrayView2<f64>) -> Array2<f64>,
        x: ArrayView2<f64>,
        initial_w: ArrayView2<f64>,
        eta: f64,
        n_iters: usize,
        epsilon: f64,
    ) -> Array2<f64> {
        let mut w = Self::direction(initial_w);
        let mut cur_iter = 0;
        while cur_iter < n_iters {
            let gradient = df(x, w.view());
            let last_w = w.clone();
            w = w + eta * gradient;
            w = Self::direction(w.view());
            if (Self::f(x, w.view()) - Self::f(x, last_w.view())).abs() < epsilon {
                break;
            }
            cur_iter += 1;
        }
        w
    }

    // 把一个w变成模为1
    fn direction(w: ArrayView2<f64>) -> Array2<f64> {
        &w / w.norm()
    }

    // 封装成求出前n个主成分的函数
    fn first_n_component(n: usize, x: ArrayView2<f64>, eta: f64, n_iters: usize) -> Array2<f64> {
        let columns = x.len_of(Axis(1));
        debug_assert!(n >= 1 && n <= columns);
        let mut res = Array2::zeros((columns, 0));
        let mut x_pca = Self::mean_to_zero(x);
        for _ in 0..n {
            let initial_w = Array2::random((x.shape()[1], 1), Uniform::new(0.0, 1.0));
            let w = Self::first_component(
                Self::df_math,
                x_pca.view(),
                initial_w.view(),
                eta,
                n_iters,
                1e-8,
            );
            res.push_column(w.to_shared().reshape(w.len_of(Axis(0))).view())
                .unwrap();
            x_pca = &x_pca - x_pca.dot(&w) * w.t();
        }
        res
    }
}

// 均值归一化
// 返回一个新的矩阵
fn mean_to_zero(x: ArrayView2<f64>) -> Array2<f64> {
    &x - x.mean_axis(Axis(0)).unwrap()
}

// 梯度上升法求解PCA
// 损失函数:f(x) = 1/m sigma(i=1,m) (XW)^2
// 梯度是: 2/m * X^T * (XW)

// 我们的目标函数
fn f(x: ArrayView2<f64>, w: ArrayView2<f64>) -> f64 {
    // 生成一个列向量
    let c1 = x.dot(&w);
    let c1 = &c1 * &c1;
    c1.sum() / x.len_of(Axis(0)) as f64
}

// 目标函数对应梯度的方法
fn df_math(x: ArrayView2<f64>, w: ArrayView2<f64>) -> Array2<f64> {
    x.t().dot(&x.dot(&w)) * 2f64 / x.len_of(Axis(0)) as f64
}

// debug函数
fn df_debug(x: ArrayView2<f64>, w: ArrayView2<f64>) -> Array2<f64> {
    let mut res: Array2<f64> = Array2::zeros((w.len_of(Axis(0)), 1));
    for i in 0..w.len_of(Axis(0)) {
        let mut w_1 = w.to_owned();
        w_1.row_mut(i)[0] += 1e-8;
        let mut w_2 = w.to_owned();
        w_2.row_mut(i)[0] -= 1e-8;
        res.row_mut(i)[0] = (f(x, w_1.view()) - f(x, w_2.view())) / (2.0 * 1e-8);
    }
    res
}

// 梯度上升法求出第一主成分
fn first_component(
    df: impl Fn(ArrayView2<f64>, ArrayView2<f64>) -> Array2<f64>,
    x: ArrayView2<f64>,
    initial_w: ArrayView2<f64>,
    eta: f64,
    n_iters: usize,
    epsilon: f64,
) -> Array2<f64> {
    let mut w = direction(initial_w);
    let mut cur_iter = 0;
    while cur_iter < n_iters {
        let gradient = df(x, w.view());
        let last_w = w.clone();
        w = w + eta * gradient;
        w = direction(w.view());
        if (f(x, w.view()) - f(x, last_w.view())).abs() < epsilon {
            break;
        }
        cur_iter += 1;
    }
    w
}

// 把一个w变成模为1
fn direction(w: ArrayView2<f64>) -> Array2<f64> {
    &w / w.norm()
}

// 封装成求出前n个主成分的函数
fn first_n_component(n: usize, x: ArrayView2<f64>, eta: f64, n_iters: usize) -> Array2<f64> {
    let columns = x.len_of(Axis(1));
    debug_assert!(n >= 1 && n <= columns);
    let mut res = Array2::zeros((columns, 0));
    let mut x_pca = mean_to_zero(x);
    for _ in 0..n {
        let initial_w = Array2::random((x.shape()[1], 1), Uniform::new(0.0, 1.0));
        let w = first_component(df_math, x_pca.view(), initial_w.view(), eta, n_iters, 1e-8);
        res.push_column(w.to_shared().reshape(w.len_of(Axis(0))).view())
            .unwrap();
        x_pca = &x_pca - x_pca.dot(&w) * w.t();
    }
    res
}
